from pyspark.sql import SparkSession

if __name__ == '__main__':
    spark = SparkSession.builder. \
        appName("test"). \
        master("local[*]"). \
        getOrCreate()

    sc = spark.sparkContext

    df = spark.read.format("csv") \
        .schema("id int, subject string, score int") \
        .load("../data/input/stu_score.txt")

    # Column对象的获取
    id_column = df["id"]
    subject_column = df["subject"]

    # DSL风格
    df.select("id", "score").show()
    df.select("id", "subject").show()
    df.select(id_column, subject_column).show()

    # filter
    df.filter("score < 99").show()
    df.filter(df["score"] < 99).show()

    # groupBy
    df.groupBy("subject").count().show()
    df.groupBy(df["subject"]).count().show()

    # df.groupBy -> GroupedData
    # 它是一个 有分组关系的数据结构，有一些API供我们对分组做聚合
    # SQL： group by 后接上聚合： sum avg count min max
    # GroupedData 类似于SQL分组后的数据结构，同样有上述5种聚合方法
    # GroupedData 是一个中转对象，调研聚合方法后，依旧返回DataFrame
    r = df.groupBy("subject")
    print(type(r))
    r.sum().show()
